#!/bin/python
"""
This python script demonstrates the MD simulation (NVE) of an FCC crystal with
Lennard-Jones potential. Especially, it demonstrates:
1) The definition of the potential;
2) The generation of the initial configuration;
3) The initialization of the temperture to a desired value;
4) The application of the periodic boundary condition;
5) The calculation of force;
6) The velocity verlet algorithm;
7) The output of the dump, therml info to files.

"""
import numpy as np

# Define the Potential parameters and the functions
sigma = 1.
epsilon = 1.
rcut  = 2.5 * sigma
rcut2 = rcut * rcut
mass  = 1.

def LJ_energy(r):
  sr = sigma / r
  sr6 = sr ** 6
  sr12 = sr6 * sr6
  return 4.*epsilon * (sr12 - sr6)

def LJ_FirstDir(r):
  sr = sigma / r
  sr6 = sr ** 6
  sr12 = sr6 * sr6
  return 24.*sr / sigma * (-2*sr12 + sr6)

# Define the simulation box, an nx x ny x nz FCC crystal, with periodic boundary condtion
alat = 1.54 * sigma   # Lattice constant
nx = 3
ny = 4
nz = 5
xbox = nx * alat
ybox = ny * alat
zbox = nz * alat
box = np.array([xbox, ybox, zbox])

def Apply_PBC(xij):
  xij = xij / box
  xij = xij - np.ceil(xij - 0.5)
  xij = xij * box
  return xij

# Initial positions FCC crystal
pos = list()

for i in range(nx):
  for j in range(ny):
    for k in range(nz):
      pos.append([i, j, k] )
      pos.append([i+0.5, j+0.5, k])
      pos.append([i+0.5, j, k+0.5])
      pos.append([i, j+0.5, k+0.5])
pos = np.array(pos) * alat
natoms = nx * ny * nz * 4;

# Initialize temperature
kb = 1.                                    # Boltzmann constant
temperature = 0.1                          # target temperature
# randomly assign velocity, within [-0.5, 0.5]
vel = np.random.random((natoms, 3)) - 0.5
vel = vel - np.sum(vel, axis=0)/natoms     # zero the momentum
kin = 0.5*mass*np.sum(vel*vel)

# rescale to get target temperature
scale = np.sqrt(1.5*natoms*kb*temperature / kin)
vel = vel * scale

# Subroutine to calculate force
def Compute_Force():
  Epot = 0.
  force = np.zeros((natoms, 3))
  for i in range(natoms):
    for j in range(i+1, natoms):
      xij = pos[j] - pos[i]

      xij = Apply_PBC(xij)

      rij2 = np.sum(xij*xij)

      if (rij2 > rcut2):
        continue

      r = np.sqrt(rij2)
      Epot += LJ_energy(r)

      fij = LJ_FirstDir(r) * xij / r

      force[i] += fij
      force[j] -= fij
  return Epot, force

# Calculate and output/write thermo info, to both file and screen
def Thermo(istep):
   Ekin = 0.5*mass*np.sum(vel*vel)
   Etot = Epot + Ekin
   T = Ekin / (1.5*natoms*kb)

   fp.write(f"{(istep+1)*dt} {T} {Epot} {Ekin} {Etot}\n")
   print("{0:.4f} {1:.6f} {2:.6f} {3:.6f} {4:.6f}".format((istep+1)*dt, T, Epot, Ekin, Etot))

# Dump atomic configuration as xyz file
def Dump(istep):
  fd.write(f"{natoms}\nStep: {istep+1}\n")
  for i in range(natoms):
    fd.write(f"X {pos[i][0]} {pos[i][1]} {pos[i][2]}\n")

# Setup
# -- control parameters
timestep = 0.002
dt = timestep
hdt = 0.5 * dt
iout  = 5
idump = 10
nsteps = 1000

# -- 0th step
Epot, force = Compute_Force()
acc = force / mass

# -- Log file
fp = open("log.dat", "w")
fp.write("# time temperature Epot Ekin Etot\n")
print("# time temperature Epot Ekin Etot")
Thermo(0)
# -- dump file
fd = open("dump.xyz", "w")
Dump(0)

# Main Loop: Velocity verlet
for istep in range(nsteps):
  vel = vel + acc * hdt
  pos = pos + vel * dt

  Epot, force = Compute_Force()
  acc = force / mass

  vel = vel + acc * hdt

  if (istep % iout  == 0):
     Thermo(istep)
  if (istep % idump == 0):
     Dump(istep)

# Clean-up
fp.close()
fd.close()
